#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import argparse
import logging
from PyPDF2 import PdfReader, PdfWriter

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def split_pdf(pdf_path, pages_per_file=7, output_dir=None, output_name_prefix=None):
    """
    将PDF文件按指定页数拆分为多个PDF文件
    
    参数:
        pdf_path: 输入PDF文件路径
        pages_per_file: 每个输出PDF文件包含的页数
        output_dir: 输出目录，默认为原PDF所在目录
        output_name_prefix: 输出文件名前缀，默认为原PDF文件名
    
    返回:
        生成的PDF文件路径列表
    """
    # 处理输出目录和文件名前缀
    if not output_dir:
        output_dir = os.path.dirname(pdf_path) or '.'
    
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    if not output_name_prefix:
        base_name = os.path.basename(pdf_path)
        output_name_prefix = os.path.splitext(base_name)[0]
    
    logger.info(f"开始拆分PDF: {pdf_path}")
    logger.info(f"每个文件 {pages_per_file} 页")
    logger.info(f"输出目录: {output_dir}")
    
    # 读取PDF文件
    try:
        reader = PdfReader(pdf_path)
        total_pages = len(reader.pages)
        logger.info(f"PDF总页数: {total_pages}")
        
        # 计算需要生成的文件数
        num_output_files = (total_pages + pages_per_file - 1) // pages_per_file
        logger.info(f"将生成 {num_output_files} 个PDF文件")
        
        output_files = []
        
        # 分割PDF
        for i in range(num_output_files):
            start_page = i * pages_per_file
            end_page = min((i + 1) * pages_per_file, total_pages)
            
            writer = PdfWriter()
            
            # 添加页面到新文件
            for page_num in range(start_page, end_page):
                writer.add_page(reader.pages[page_num])
            
            # 生成输出文件名
            output_filename = f"{output_name_prefix}_part{i+1}.pdf"
            output_path = os.path.join(output_dir, output_filename)
            
            # 保存拆分后的PDF
            with open(output_path, 'wb') as output_file:
                writer.write(output_file)
            
            output_files.append(output_path)
            logger.info(f"已生成: {output_path} (页码 {start_page+1}-{end_page})")
        
        logger.info(f"PDF拆分完成，共生成 {len(output_files)} 个文件")
        return output_files
    
    except Exception as e:
        logger.error(f"拆分PDF时发生错误: {str(e)}")
        logger.exception(e)
        return []

def main():
    parser = argparse.ArgumentParser(description='将PDF文件按固定页数拆分为多个小文件')
    parser.add_argument('pdf_path', help='要拆分的PDF文件路径')
    parser.add_argument('-n', '--pages', type=int, default=20, help='每个输出PDF包含的页数（默认为7）')
    parser.add_argument('-o', '--output-dir', help='输出目录（默认为原PDF所在目录）')
    parser.add_argument('-p', '--prefix', help='输出文件名前缀（默认为原PDF文件名）')
    
    args = parser.parse_args()
    
    split_pdf(args.pdf_path, args.pages, args.output_dir, args.prefix)

if __name__ == "__main__":
    main()
